#-*- coding:utf-8 -*-
# author:贤宁
# datetime:2021/12/2 15:59
# software: PyCharm

import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from utils import get_conv_layer
from torch.distributions import uniform
"""三种残差块"""
#G Block
class G_Block(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(G_Block,self).__init__()
        self.block1 = nn.Sequential(
            nn.UpsamplingBilinear2d(scale_factor=2),  #线性上采样
            nn.Conv2d(in_channel, out_channel, 1),   #1x1 卷积
        )

        self.block2 = nn.Sequential(
            nn.BatchNorm2d(in_channel),
            nn.LeakyReLU(0.2,inplace=True),
            nn.UpsamplingBilinear2d(scale_factor=2),   #线性上采样
            nn.Conv2d(in_channel, in_channel, 3, stride=1, padding=1),   #这里特征图与通道数都不变
            nn.BatchNorm2d(in_channel),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(in_channel, out_channel, 3, stride=1, padding=1)   #这里特征图大小不变，通道数增加
        )
    def forward(self,x):
        return self.block1(x) + self.block2(x)   #残差求和
#D&3D Block
class D_Block(nn.Module):
    def __init__(self,in_channel,out_channel,conv_type,first_relu=True,keep_same_output=False):
        """keep_same_output：输出是否与输入具有相同的空间维度，若为FALSE，则缩小为2
           conv_type: 决定卷积的类型,"standard"--->Conv2d,"3d"---->Conv3d
           first_relu:也为一个bool型
        """
        super(D_Block, self).__init__()
        self.keep_same_output = keep_same_output
        self.conv_type = conv_type
        self.first_relu = first_relu
        conv2d = get_conv_layer(conv_type)
        self.conv_1x1 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel,
            kernel_size=1,
        )
        self.first_conv_3x3 = conv2d(
            in_channels=in_channel,
            out_channels=in_channel,
            kernel_size=3,
            padding=1,
        )
        self.last_conv_3x3 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel,
            kernel_size=3,
            padding=1,
            stride=1,
        )
        if conv_type == "3d":
            # 需要谱的归一化
            self.conv_1x1 = spectral_norm(self.conv_1x1)
            self.first_conv_3x3 = spectral_norm(self.first_conv_3x3)
            self.last_conv_3x3 = spectral_norm(self.last_conv_3x3)
        self.relu = torch.nn.LeakyReLU(0.2,inplace=True)
    def forward(self,x):
        x1 = self.conv_1x1(x)
        if not self.keep_same_output:
            x1 = F.interpolate(
                x1, mode="trilinear" if self.conv_type == "3d" else "bilinear", scale_factor=0.5
            )
        if self.first_relu:
            x = self.relu(x)
        x = self.first_conv_3x3(x)
        x = self.relu(x)
        x = self.last_conv_3x3(x)
        if not self.keep_same_output:
            x = F.interpolate(
                x, mode="trilinear" if self.conv_type == "3d" else "bilinear", scale_factor=0.5
            )
        x = x1 + x  # 计算残差
        return x
#L Block(整体均不会改变特征图的大小)
class L_Block(nn.Module):
    def __init__(self,in_channel,out_channel, conv_type="standard"):
        super(L_Block, self).__init__()
        self.conv_type = conv_type
        conv2d = get_conv_layer(conv_type)
        self.conv_1x1 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel - in_channel,
            kernel_size=1,
        )

        self.first_conv_3x3 = conv2d(
            in_channels= in_channel, out_channels=out_channel, kernel_size=3, padding=1, stride=1
        )
        self.relu = torch.nn.LeakyReLU(0.2,inplace=True)
        self.last_conv_3x3 = conv2d(
            in_channels=out_channel,
            out_channels=out_channel,
            kernel_size=3,
            padding=1,
            stride=1,
        )
    def forward(self,x):
        x1 = self.conv_1x1(x)
        #print(x1.size())
        x = self.relu(x)
        x2 = self.first_conv_3x3(x)
        x2 = self.relu(x2)
        x2 = self.last_conv_3x3(x2)
        #print(x2.size())
        x = x2 + (torch.cat((x,x1),dim=1))
        return x
"""Block---End"""